package test.demo;

import com.alibaba.fastjson.JSON;
import tech.waterism.modelbase.PredictBase;
import tech.waterism.modelbase.PredictResult;

import java.util.Arrays;
import java.util.Calendar;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 陕北模型 超渗产流
 *
 * @author 初京刚
 * @version 1.0
 * @date 2022-3-4
 */

public class SHANNXI_R extends PredictBase {

    /**
     * 流域平均张力水容量（mm）、流域蓄水容量
     */
    private double Wm;

    /**
     * 流域蒸散发折算系数
     */
    private double Kc;

    /**
     * 流域内下渗能力抛物线形状系数
     */
    private double Bx;

    /**
     * 不透水的面积占全流域面积的比
     */
    private double IM;

    /**
     * 流域 月蒸发量（mm）
     */
    private double[] ES;

    /**
     * 流域初始下渗能力，相当于土壤干燥时的下渗能力，平均最大下渗能力，mm/min
     */
    private double f0;

    /**
     * 流域平均稳定下渗能力，mm/min
     */
    private double fc;

    /**
     * 土壤含水计算的允许误差
     */
    private double err_limit;

    /**
     * 下渗能力衰弱系数，随土质而变的系数，1/min
     */
    private double k;

    /**
     * 流域 时段蒸发量（mm）
     */
    private double[] floodE;

    /**
     * 降水时间
     */
    private Date[] floodTm;

    /**
     * 降水量
     */
    private double[] floodDrp;

    /**
     * 降水序列长度
     */
    private Integer floodRainRange;

    /**
     * 降水时段长
     */
    private double timeInterval;

    /**
     * 模拟总产流
     */
    private double[] RunoffSim;

    /**
     * 总产流
     */
    private double[] R;

    //    /**
    //     * 不透水面积产流
    //     */
    //    private double[] RIm;

    /**
     * 净雨
     */
    private double[] PE;

    /**
     * 雨强 mm/min
     */
    private double[] PEI;

    /**
     * 张力水含量
     */
    private double[] Ws;

    /**
     * 流域面积
     */
    private double Area;

    /**
     * 流域张力水含量初值（m）
     */
    private double WSInitial;

    /**
     * 下渗曲线选取标志
     */
    private int type;

    /**
     * 菲利普下渗曲线参数
     */
    private double A;

    /**
     * 菲利普下渗曲线参数
     */
    private double B;

    public SHANNXI_R(Map<String, String> data) {
        super();
        WSInitial = Double.parseDouble(data.get("WSIG"));
        Kc = Double.parseDouble(data.get("KC"));
        Bx = Double.parseDouble(data.get("BX"));
        IM = Double.parseDouble(data.get("IM"));
        type = Integer.parseInt(data.get("TYPE"));
        f0 = Double.parseDouble(data.get("F0"));
        fc = Double.parseDouble(data.get("FC"));
        err_limit = Double.parseDouble(data.get("ERLMT"));
        k = Double.parseDouble(data.get("K"));
        Wm = f0 / k;
        A = Double.parseDouble(data.get("A"));
        B = Double.parseDouble(data.get("B"));
        ES = JSON.parseObject(data.get("ES"), double[].class);
        floodTm = JSON.parseObject(data.get("dt"), Date[].class);// input[i].data.dt
        floodDrp = JSON.parseObject(data.get("rain"), double[].class);// input[i].data.rain
        floodE = JSON.parseObject(data.get("evaporation"), double[].class);  // input[i].data.evaporation
        floodRainRange = floodDrp.length;
        timeInterval = Double.parseDouble(data.get("clen"));                  // index.clen
        RunoffSim = new double[floodRainRange];
        R = new double[floodRainRange];
        //        RIm = new double[floodRainRange];
        PE = new double[floodRainRange];
        PEI = new double[floodRainRange];
    }

    /**
     * 模型计算
     */
    public Map<String, double[]> SHANNXImodel() {

        // 超渗产流
        expinfMechanism();

        for (Integer j = 0; j < floodRainRange; j++) {
            //            RunoffSim[j] = RS[j] + RIm[j] + RI[j] + RG[j];
            RunoffSim[j] = R[j];
        }

        Map<String, double[]> simulationResult = new HashMap<>();
        simulationResult.put("rain", floodDrp);
        simulationResult.put("PE", PE);
        simulationResult.put("runoffSim", RunoffSim);
        simulationResult.put("R", R);
        //        simulationResult.put("RIM", RIm);

        return simulationResult;
    }

    /**
     * 蒸发计算
     *
     * @param date   降水时间
     * @param timeIn 时间间隔 (h)
     * @return 蒸发
     */
    private double calcEvaporationPotential(Date date, double timeIn) {
        Calendar time = Calendar.getInstance();
        time.setTime(date);
        Integer Month = time.get(Calendar.MONTH) + 1; // 获得降雨开始月份
        double Ep = ES[Month - 1] / 30; // 月蒸发量转换成日蒸发量
        Ep = Ep * timeIn / 24; // 日蒸发量转换成时段蒸发量
        return Ep;
    }

    /**
     * 超渗产流模型
     *
     * @return 总产流
     */
    private void expinfMechanism() {

        Ws = new double[floodRainRange + 1]; // 张力水蓄量，随时段变化
        Ws[0] = WSInitial;

        double EPt; // EPt为蒸散发，等于Kc * E（实测蒸发，即ES[]或时段）
        double Ft0 = 0.0; // 时段初流域下渗能力
        double Ft = 0.0; //时段下渗量
        double fmm;

        for (Integer i = 0; i < floodRainRange; i++) {

            //蒸发计算
            if (floodE == null || floodE.length != floodRainRange) {
                EPt = Kc * calcEvaporationPotential(floodTm[i], timeInterval);
            } else {
                EPt = Kc * floodE[i];
            }

            // PE[i]为降雨减去蒸发，可正可负
            PE[i] = floodDrp[i] - EPt;
            PEI[i] = PE[i] / (60 * timeInterval);

            if (type == 0) {
                //霍顿下渗曲线<霍顿与菲利蒲下渗公式对子洲径流站资料的拟合>
                Ft0 = HortonInfil(Ws[i]);
                // 计算流域该时段的最大点下渗能力<格林一安普特下渗曲线的改进和应用>、<格林－安普特降雨径流模型改进及初步应用>
                fmm = Ft0 * (1 + Bx) / (1 - IM);
                if (PEI[i] > fmm) {
                    // 全流域产流
                   /* Ws[i + 1] = fc * 60 * timeInterval * (i + 1) + 1 / k * (f0 - fc) * (1 - Math.exp(-k * 60 * timeInterval * (i + 1)));
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                    Ft = Ws[i + 1] - Ws[i];*/
                    Ft = Ft0 * 60 * timeInterval;
                    Ws[i + 1] = Ws[i] + Ft;
                    if(Ft < 0){
                        Ft = 0;
                    }
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                    if(Ws[i + 1] < 0){
                        Ws[i + 1] = 0;
                    }
                    R[i] = PE[i] - Ft;
                    if(R[i] < 0){
                        R[i] = 0;
                    }
                } else {
                    R[i] = (PEI[i] - Ft0 + Ft0 * Math.pow(1 - PEI[i] / fmm, Bx + 1)) * 60 * timeInterval;
                    if(R[i] < 0){
                        R[i] = 0;
                    }
                    Ws[i + 1] = Ws[i] + PE[i] - R[i];
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                    if(Ws[i + 1] < 0){
                        Ws[i + 1] = 0;
                    }
                }
            } else if (this.type == 1) {
                //菲利普斯下渗曲线<霍顿与菲利蒲下渗公式对子洲径流站资料的拟合>
                if(Ws[i] > 0) {
                    Ft0 = PhillipsInfli(Ws[i]);
                } else {
                    Ft0 = PhillipsInfli(0.1);
                }
                // 计算流域该时段的最大点下渗能力<格林一安普特下渗曲线的改进和应用>、<格林－安普特降雨径流模型改进及初步应用>
                fmm = Ft0 * (1 + Bx) / (1 - IM);
                if (PEI[i] > fmm) {
                    // 全流域产流
                    R[i] = (PEI[i] - Ft0) * 60 * timeInterval;
                    if(R[i] < 0){
                        R[i] = 0;
                    }
                    Ws[i + 1] = Ws[i] + Ft0 * 60 * timeInterval;
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                } else {
                    R[i] = (PEI[i] - Ft0 + Ft0 * Math.pow(1 - PEI[i] / fmm, Bx + 1)) * 60 * timeInterval;
                    if(R[i] < 0){
                        R[i] = 0;
                    }
                    Ws[i + 1] = Ws[i] + PE[i] - R[i];
                    if(Ws[i + 1] > Wm){
                        Ws[i + 1] = Wm;
                    }
                    if(Ws[i + 1] < 0){
                        Ws[i + 1] = 0;
                    }
                }
            }
        }
    }

    /**
     * 菲利普斯下渗曲线计算
     *
     * @param w0 时段初土壤含水
     * @return 下渗能力
     */
    private double PhillipsInfli(double w0) {
        return B * B * (1 + Math.sqrt(1 + 4 * A * w0 / (B * B))) / (4 * w0) + A;
    }

    /**
     * 霍顿下渗曲线计算
     *
     * @param w0 时段初土壤含水
     * @return 下渗能力
     */
    private double HortonInfil(double w0) {

        /*
            理论基础：<流域水文模型-新安江模型和陕北模型>pdf66页
            数学基础：迭代求解
        */
        // 由经验初设模型参数
        double T = w0 / f0;

        double w0ByCount = fc * T + 1 / k * (f0 - fc) * (1 - Math.exp(-k * T));
        double fByCount = 0.0;
        while (Math.abs(w0 - w0ByCount) > err_limit) {
            fByCount = f0 - k * (w0ByCount - fc * T);
//            fByCount = fc + (f0 - fc) * Math.exp(-k * T);
            T = T + (w0 - w0ByCount) / fByCount;
            w0ByCount = fc * T + 1 / k * (f0 - fc) * (1 - Math.exp(-k * T));
        }
        fByCount = f0 - k * (w0 - fc * T);
        return fByCount;
    }

    private List<String> toStringList(double[] array0) {
        String[] array = new String[array0.length];
        for (int i = 0; i < array.length; i++) {
            array[i] = String.valueOf(array0[i]);
        }
        return Arrays.asList(array);
    }

    private List<String> toStringList(double[][] array0) {
        String[] array = new String[array0.length];
        for (int i = 0; i < array.length; i++) {
            array[i] = Arrays.toString(array0[i]);
        }
        return Arrays.asList(array);
    }

    @Override
    public void init() {

    }

    @Override
    public PredictResult predict() {
        PredictResult result = new PredictResult();
        Map<String, double[]> map = SHANNXImodel();

        double[][] RSim = new double[1][];
        //总产流
        RSim[0] = map.get("R");
        //        //流域地表产流
        //        RSim[1] = map.get("RS");
        //        //流域土壤产流
        //        RSim[2] = map.get("RI");
        //        //流域地下产流
        //        RSim[3] = map.get("RG");

        //流域面积总产流
        result.setRunoffSim(map.get("runoffSim"));
        result.setRSim(RSim);

        // 更新状态
        List<String> WS = toStringList(Ws);
        //        List<String> WUP = toStringList(Wu);
        //        List<String> WLP = toStringList(Wl);
        //        List<String> WDP = toStringList(Wd);
        //        List<String> SP = toStringList(S);
        //        List<String> FRP = toStringList(FR);
        //        List<String> QXSIG = toStringList(QX0);
        //        List<String> QSIG = toStringList(QS0);
        //        List<String> QSP = toStringList(QS);
        //        List<String> QIP = toStringList(QIS0);
        //        List<String> QGP = toStringList(QGS0);
        //        List<String> QSP = toStringList(QS);
        //        List<String> QIP = toStringList(QI);
        //        List<String> QGP = toStringList(QG);
        Map<String, List<String>> status = new HashMap<>();
        //土壤含水量
        status.put("WS", WS);
        //        //上层张力水含量
        //        status.put("WUP", WUP);
        //        //下层张力水含量
        //        status.put("WLP", WLP);
        //        //深层张力水含量
        //        status.put("WDP", WDP);
        //        //产流面积自由水深
        //        status.put("SP", SP);
        //        //产流面积系数
        //        status.put("FRP", FRP);
        //        //河段下断面各时段初始流量
        //        status.put("QXSIG", QXSIG);
        //        //各时段初始总基流
        //        status.put("QSIG", QSIG);
        //        //地表坡面汇流
        //        status.put("QSP", QSP);
        //        //坡面汇流-壤中流
        //        status.put("QIP", QIP);
        //        //地下坡面汇流
        //        status.put("QGP", QGP);
        result.setNewStatus(status);

        return result;

        //return null;
    }
}
